package com.czk.tree.huffmantree.huffmancode;

import lombok.Data;

import java.util.*;

/**
 * 赫夫曼编码、压缩的代码
 *
 * @Author:ChenZhangKun
 * @Date: 2020/7/10 22:25
 */
public class HuffmanCode {
    static Map<Byte, String> huffmanMap = new HashMap<>();
    static StringBuilder builder = new StringBuilder();

    public static void main(String[] args) {
        String str = "i like like like java do you like a java";
        // 拿到字数数组 长度40
        byte[] bytes = str.getBytes();
        System.out.println(Arrays.toString(bytes));
        // 得到字符串对应的list集合
        List<Node> node = getNode(bytes);
        // 创建赫夫曼树
        Node huffmanTree = createHuffmanTree(node);
        preTraverse(huffmanTree);
        System.out.println(node);
        // 生成哈夫曼编码表
        Map<Byte, String> huffmanTable = createHuffmanTable(huffmanTree);
        System.out.println("=============");
        System.out.println(huffmanTable);
        System.out.println("赫夫曼编码的二进制");
        byte[] zip = zip(bytes, huffmanMap);
        System.out.println(Arrays.toString(zip));// 长度为17，原始长度为40
        System.out.println(Arrays.toString(huffmanZip(bytes)));
        System.out.println("转码成二进制");

        byte[] decode = decode(huffmanMap, zip);
        System.out.println(new String(decode));

    }

    /**
     * @param huffmanMap   赫夫曼编码对应的二进制字符串
     * @param huffmanBytes 压缩后形成的数组
     * @return
     */
    public static byte[] decode(Map<Byte, String> huffmanMap, byte[] huffmanBytes) {
        StringBuilder builder = new StringBuilder();
        // 得到Huffman字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            // 判断是不是最后一个字节
            boolean flag = (i == huffmanBytes.length - 1);
            // 是最后一个字节不需要补码
            String string = byteToBitString(!flag, huffmanBytes[i]);
            builder.append(string);
        }
        System.out.println(builder.toString());
        // 把Huffman编码表进行对调
        Map<String, Byte> map = new HashMap<>();
        // 遍历HuffmanMap进行对调
        for (Byte aByte : huffmanMap.keySet()) {
            map.put(huffmanMap.get(aByte), aByte);
        }
        // 遍历集合，存放byte到list集合中
        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < builder.length(); ) {
            int count = 1;
            boolean flag = true;
            Byte b = null;
            while (flag) {
                String substring = builder.substring(i, i + count);
                b = map.get(substring);
                if (b == null) {
                    count++;
                } else {
                    flag = false;
                }
            }
            list.add(b);
            i += count;
        }
        // 将list中的数据放入集合中并返回
        byte[] bytes = new byte[list.size()];
        for (int i = 0; i < list.size(); i++) {
            bytes[i] = list.get(i);
        }
        System.out.println(Arrays.toString(bytes));
        return bytes;
    }

    /**
     * 数据的解压，先将byte数组转化成二进制数据·
     *
     * @param bytes 需要解码的字节
     * @return
     */
    public static String byteToBitString(boolean flag, byte bytes) {
        // 使用变量保存byte
        int temp = bytes;
        // 如果是正数，我们进行补高位
        if (flag) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);// 返回temp对应的二进制补码
        // 如果是最后一个字节不需要补高位
        if (flag) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }

    /**
     * 得到经过哈夫曼编码处理后的数组
     *
     * @param originBytes 原始字符串的字符数组
     * @return
     */
    public static byte[] huffmanZip(byte[] originBytes) {
        // 得到对应的list<Node集合>
        List<Node> node = getNode(originBytes);
        // 创建赫夫曼树
        Node huffmanTree = createHuffmanTree(node);
        // 生成赫夫曼编码表
        huffmanMap = createHuffmanTable(huffmanTree);
        // 生成压缩
        byte[] zip = zip(originBytes, huffmanMap);
        return zip;
    }

    /**
     * 返回赫夫曼编码处理后的编码数组
     *
     * @param bytes      原始的byte数组
     * @param huffmanMap 赫夫曼编码表
     * @return
     */
    public static byte[] zip(byte[] bytes, Map<Byte, String> huffmanMap) {
        // 利用huffmanMap将原始byte数组转换成010101010001的二进制数据
        StringBuilder builder = new StringBuilder();
        for (byte aByte : bytes) {
            builder.append(huffmanMap.get(aByte));
        }
        // builder此时是对应的二进制字符串
        System.out.println(builder.toString() + "  " + builder.toString().length());
        // 将对应的字符串装成byte数组
        // 能整除则刚好
        int len;
        if (builder.length() % 8 == 0) {
            len = builder.length() / 8;
        } else {
            // 不能整除，多加一个
            len = builder.length() / 8 + 1;
        }
        // 创建数组
        byte[] bytes1 = new byte[len];
        // 记录第几个byte
        int index = 0;
        // 步长为8
        for (int i = 0; i < builder.length(); i += 8) {
            // 截取字符串
            String stringByte;
            if (i + 8 > builder.length()) {
                stringByte = builder.substring(i);
            } else {
                stringByte = builder.substring(i, i + 8);
            }
            // 将stringByte二进制转成byte，放在数组中
            bytes1[index] = (byte) Integer.parseInt(stringByte, 2);
            index++;
        }
        return bytes1;
    }

    /**
     * 重载方法，封装参数
     *
     * @param huffmanTreeNode
     * @return
     */
    public static Map<Byte, String> createHuffmanTable(Node huffmanTreeNode) {
        if (huffmanTreeNode == null) {
            return null;
        }
        return createHuffmanTable(huffmanTreeNode, "", builder);
    }

    /**
     * 便于外部拿到huffmanMap
     *
     * @return
     */
    public static Map<Byte, String> getHuffmanMap() {
        return huffmanMap;
    }

    /**
     * 创建对应的赫夫曼编码表
     *
     * @param huffmanTreeNode 根节点
     * @param code            路径 左子节点为0，右子节点为1
     * @param builder         拼接路径
     * @return
     */
    private static Map<Byte, String> createHuffmanTable(Node huffmanTreeNode, String code, StringBuilder builder) {
        StringBuilder stringBuilder = new StringBuilder(builder);
        stringBuilder.append(code);
        if (huffmanTreeNode != null) {
            // 判断当前节点是叶子节点还是非叶子节点
            if (huffmanTreeNode.data == null) {
                // 非叶子节点
                // 向左递归
                createHuffmanTable(huffmanTreeNode.left, "0", stringBuilder);
                // 向右递归
                createHuffmanTable(huffmanTreeNode.right, "1", stringBuilder);
            } else {
                // 叶子节点,加入map中
                huffmanMap.put(huffmanTreeNode.data, stringBuilder.toString());
            }

        }
        return huffmanMap;
    }

    /**
     * 前序遍历
     *
     * @param node
     */
    public static void preTraverse(Node node) {
        node.preTraverse();
    }

    /**
     * 得到赫夫曼树
     *
     * @param list
     * @return
     */
    public static Node createHuffmanTree(List<Node> list) {
        //
        while (list.size() > 1) {
            Collections.sort(list);
            Node left = list.get(0);
            Node right = list.get(1);
            // 创建一颗新的二叉树
            // 新的节点只有权值，没有data
            Node parent = new Node(null, left.weight + right.weight);
            parent.left = left;
            parent.right = right;
            list.remove(left);
            list.remove(right);
            list.add(parent);
        }
        return list.get(0);
    }

    /**
     * 得到list的Node集合
     *
     * @param bytes
     * @return
     */
    public static List<Node> getNode(byte[] bytes) {
        List<Node> list = new ArrayList<>();
        // 存储每个byte出现的次数
        Map<Byte, Integer> map = new HashMap<>();
        // 遍历bytes
        for (byte aByte : bytes) {
            if (!map.containsKey(aByte)) {
                map.put(aByte, 1);
            } else {
                Integer integer = map.get(aByte);
                integer++;
                map.put(aByte, integer);
            }
        }
        // 遍历map，存储节点
        for (Byte aByte : map.keySet()) {
            list.add(new Node(aByte, map.get(aByte)));
        }
        return list;
    }


}

@Data
class Node implements Comparable<Node> {
    Byte data; // 数据 97->a
    int weight; // 权值，字符出现的次数
    Node left; // 左节点
    Node right; //右节点

    public Node(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    // 比较器，自然顺序排序
    @Override
    public int compareTo(Node o) {
        return this.weight - o.weight;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    // 遍历
    public void preTraverse() {
        System.out.println(this);
        if (this.left != null) {
            this.left.preTraverse();
        }
        if (this.right != null) {
            this.right.preTraverse();
        }
    }
}
